import sys
import os

sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from data_module import CellDataModule
from model_module import LitMonai2DModel
import pytorch_lightning as pl
from pytorch_lightning.callbacks import ModelCheckpoint

# 导入自定义模块
from utils.swanlab_logger import SwanLabLogger
from utils.config_manager import ConfigManager


def main():
    # 1. 解析命令行参数
    args = ConfigManager.parse_args()

    # 2. 加载配置
    config_manager = ConfigManager(args.config)
    config = config_manager.update_config_from_args(args)

    # 3. 初始化数据模块
    data_module = CellDataModule(config)
    data_module.setup()
    
    # 4. 初始化模型模块
    model = LitMonai2DModel(config)

    # 5. 配置回调
    checkpoint_callback = ModelCheckpoint(
        dirpath=config["callbacks"]["checkpoint"]["dirpath"],
        filename=config["callbacks"]["checkpoint"]["filename"],
        save_top_k=config["callbacks"]["checkpoint"]["save_top_k"],
        save_last=config["callbacks"]["checkpoint"]["save_last"],
        monitor=config["callbacks"]["checkpoint"][
            "monitor"
        ],  # 监控验证指标，其传参值和model的self.log()里的标题有关系
        mode=config["callbacks"]["checkpoint"]["mode"],
    )  # 怎么理解以及使用

    # 6. 初始化 SwanLab Logger
    swanlab_logger = SwanLabLogger(
        project=config["logging"]["project"],
        experiment_name=config["logging"]["experiment_name"],
        description=config["logging"]["description"],
        logdir=config["logging"]["logdir"],
    )

    # 7. 初始化 Trainer  # 可以尝试封装为一个函数
    trainer = pl.Trainer(
        max_epochs=config["training"]["max_epochs"],
        accelerator=config["training"]["accelerator"],
        devices=config["training"]["devices"],
        callbacks=[checkpoint_callback],
        logger=swanlab_logger,
    )

    # 8. 开始训练
    trainer.fit(
        model, datamodule=data_module
    )

    # trainer.test(model, datamodule=data_module)


if __name__ == "__main__":
    main()
